package com.mini.grpc.protoc.generator;

import com.mini.grpc.protoc.model.MethodDefinition;
import com.mini.grpc.protoc.model.ServiceDefinition;

import java.io.IOException;
import java.io.Writer;
import java.util.List;

/**
 * 存根代码生成器
 * 生成gRPC服务的Java存根代码
 * 
 * @author Mini-gRPC
 */
public class StubGenerator {
    
    private static final String INDENT = "    ";
    private static final String DOUBLE_INDENT = "        ";
    private static final String TRIPLE_INDENT = "            ";
    
    /**
     * 生成服务存根代码
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    public void generateStub(ServiceDefinition service, Writer writer) throws IOException {
        generateFileHeader(service, writer);
        generateServiceClass(service, writer);
        generateServiceInterface(service, writer);
        generateBlockingStub(service, writer);
        generateAsyncStub(service, writer);
        generateFutureStub(service, writer);
        generateMethodDescriptors(service, writer);
        generateServiceDescriptor(service, writer);
        generateFactoryMethods(service, writer);
        generateFileFooter(writer);
    }
    
    /**
     * 生成文件头部
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateFileHeader(ServiceDefinition service, Writer writer) throws IOException {
        if (service.getJavaPackage() != null) {
            writer.write("package " + service.getJavaPackage() + ";\n\n");
        }
        
        // 生成导入语句
        writer.write("import com.mini.grpc.client.AbstractStub;\n");
        writer.write("import com.mini.grpc.client.MiniGrpcChannel;\n");
        writer.write("import com.mini.grpc.client.ClientCall;\n");
        writer.write("import com.mini.grpc.client.CallOptions;\n");
        writer.write("import com.mini.grpc.common.MethodDescriptor;\n");
        writer.write("import com.mini.grpc.common.Status;\n");
        writer.write("import com.mini.grpc.server.BindableService;\n");
        writer.write("import com.mini.grpc.server.ServerServiceDefinition;\n");
        writer.write("import com.mini.grpc.server.ServerMethodDefinition;\n");
        writer.write("import com.mini.grpc.server.ServerCallHandler;\n");
        writer.write("import com.mini.grpc.stream.StreamObserver;\n");
        writer.write("import com.mini.grpc.serialization.Marshaller;\n");
        writer.write("import com.mini.grpc.serialization.json.JsonMarshaller;\n");
        writer.write("\nimport java.util.concurrent.Future;\n");
        writer.write("import java.util.concurrent.CompletableFuture;\n\n");
    }
    
    /**
     * 生成服务类
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateServiceClass(ServiceDefinition service, Writer writer) throws IOException {
        writer.write("/**\n");
        writer.write(" * " + service.getServiceName() + " gRPC service\n");
        writer.write(" * Generated by Mini-gRPC protoc plugin\n");
        writer.write(" */\n");
        writer.write("public final class " + service.getJavaClassName() + " {\n\n");
        
        writer.write(INDENT + "private " + service.getJavaClassName() + "() {}\n\n");
        
        writer.write(INDENT + "public static final String SERVICE_NAME = \"" + 
                    service.getFullServiceName() + "\";\n\n");
    }
    
    /**
     * 生成服务接口
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateServiceInterface(ServiceDefinition service, Writer writer) throws IOException {
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * " + service.getServiceName() + " service interface\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public interface " + service.getServiceInterfaceName() + " extends BindableService {\n\n");
        
        // 生成接口方法
        for (MethodDefinition method : service.getMethods()) {
            generateInterfaceMethod(method, writer);
        }
        
        writer.write(INDENT + "}\n\n");
    }
    
    /**
     * 生成接口方法
     * 
     * @param method 方法定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateInterfaceMethod(MethodDefinition method, Writer writer) throws IOException {
        writer.write(DOUBLE_INDENT + "/**\n");
        writer.write(DOUBLE_INDENT + " * " + method.getMethodName() + " method\n");
        writer.write(DOUBLE_INDENT + " */\n");
        
        switch (method.getMethodType()) {
            case UNARY:
                writer.write(DOUBLE_INDENT + "void " + method.getJavaMethodName() + 
                           "(" + method.getJavaInputType() + " request, " +
                           "StreamObserver<" + method.getJavaOutputType() + "> responseObserver);\n\n");
                break;
            case CLIENT_STREAMING:
                writer.write(DOUBLE_INDENT + "StreamObserver<" + method.getJavaInputType() + "> " + 
                           method.getJavaMethodName() + 
                           "(StreamObserver<" + method.getJavaOutputType() + "> responseObserver);\n\n");
                break;
            case SERVER_STREAMING:
                writer.write(DOUBLE_INDENT + "void " + method.getJavaMethodName() + 
                           "(" + method.getJavaInputType() + " request, " +
                           "StreamObserver<" + method.getJavaOutputType() + "> responseObserver);\n\n");
                break;
            case BIDI_STREAMING:
                writer.write(DOUBLE_INDENT + "StreamObserver<" + method.getJavaInputType() + "> " + 
                           method.getJavaMethodName() + 
                           "(StreamObserver<" + method.getJavaOutputType() + "> responseObserver);\n\n");
                break;
        }
    }
    
    /**
     * 生成阻塞存根
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateBlockingStub(ServiceDefinition service, Writer writer) throws IOException {
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Blocking stub for " + service.getServiceName() + "\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static final class " + service.getBlockingStubClassName() + 
                    " extends AbstractStub<" + service.getBlockingStubClassName() + "> {\n\n");
        
        // 构造函数
        writer.write(DOUBLE_INDENT + "private " + service.getBlockingStubClassName() + 
                    "(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "super(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // build方法
        writer.write(DOUBLE_INDENT + "@Override\n");
        writer.write(DOUBLE_INDENT + "protected " + service.getBlockingStubClassName() + 
                    " build(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "return new " + service.getBlockingStubClassName() + 
                    "(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // 生成阻塞方法
        for (MethodDefinition method : service.getMethods()) {
            if (method.isUnary()) {
                generateBlockingMethod(method, writer);
            }
        }
        
        writer.write(INDENT + "}\n\n");
    }
    
    /**
     * 生成阻塞方法
     * 
     * @param method 方法定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateBlockingMethod(MethodDefinition method, Writer writer) throws IOException {
        writer.write(DOUBLE_INDENT + "/**\n");
        writer.write(DOUBLE_INDENT + " * " + method.getMethodName() + " blocking call\n");
        writer.write(DOUBLE_INDENT + " */\n");
        writer.write(DOUBLE_INDENT + "public " + method.getJavaOutputType() + " " + 
                    method.getJavaMethodName() + "(" + method.getJavaInputType() + " request) {\n");
        writer.write(TRIPLE_INDENT + "return blockingUnaryCall(\n");
        writer.write(TRIPLE_INDENT + INDENT + "getChannel(), " + method.getMethodDescriptorFieldName() + ", getCallOptions(), request);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
    }
    
    /**
     * 生成异步存根
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateAsyncStub(ServiceDefinition service, Writer writer) throws IOException {
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Async stub for " + service.getServiceName() + "\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static final class " + service.getAsyncStubClassName() + 
                    " extends AbstractStub<" + service.getAsyncStubClassName() + "> {\n\n");
        
        // 构造函数
        writer.write(DOUBLE_INDENT + "private " + service.getAsyncStubClassName() + 
                    "(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "super(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // build方法
        writer.write(DOUBLE_INDENT + "@Override\n");
        writer.write(DOUBLE_INDENT + "protected " + service.getAsyncStubClassName() + 
                    " build(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "return new " + service.getAsyncStubClassName() + 
                    "(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // 生成异步方法
        for (MethodDefinition method : service.getMethods()) {
            generateAsyncMethod(method, writer);
        }
        
        writer.write(INDENT + "}\n\n");
    }
    
    /**
     * 生成异步方法
     * 
     * @param method 方法定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateAsyncMethod(MethodDefinition method, Writer writer) throws IOException {
        writer.write(DOUBLE_INDENT + "/**\n");
        writer.write(DOUBLE_INDENT + " * " + method.getMethodName() + " async call\n");
        writer.write(DOUBLE_INDENT + " */\n");
        
        switch (method.getMethodType()) {
            case UNARY:
                writer.write(DOUBLE_INDENT + "public void " + method.getJavaMethodName() + 
                           "(" + method.getJavaInputType() + " request, " +
                           "StreamObserver<" + method.getJavaOutputType() + "> responseObserver) {\n");
                writer.write(TRIPLE_INDENT + "asyncUnaryCall(\n");
                writer.write(TRIPLE_INDENT + INDENT + "getChannel().newCall(" + method.getMethodDescriptorFieldName() + ", getCallOptions()),\n");
                writer.write(TRIPLE_INDENT + INDENT + "request, responseObserver);\n");
                break;
            case CLIENT_STREAMING:
                writer.write(DOUBLE_INDENT + "public StreamObserver<" + method.getJavaInputType() + "> " + 
                           method.getJavaMethodName() + 
                           "(StreamObserver<" + method.getJavaOutputType() + "> responseObserver) {\n");
                writer.write(TRIPLE_INDENT + "return asyncClientStreamingCall(\n");
                writer.write(TRIPLE_INDENT + INDENT + "getChannel().newCall(" + method.getMethodDescriptorFieldName() + ", getCallOptions()),\n");
                writer.write(TRIPLE_INDENT + INDENT + "responseObserver);\n");
                break;
            case SERVER_STREAMING:
                writer.write(DOUBLE_INDENT + "public void " + method.getJavaMethodName() + 
                           "(" + method.getJavaInputType() + " request, " +
                           "StreamObserver<" + method.getJavaOutputType() + "> responseObserver) {\n");
                writer.write(TRIPLE_INDENT + "asyncServerStreamingCall(\n");
                writer.write(TRIPLE_INDENT + INDENT + "getChannel().newCall(" + method.getMethodDescriptorFieldName() + ", getCallOptions()),\n");
                writer.write(TRIPLE_INDENT + INDENT + "request, responseObserver);\n");
                break;
            case BIDI_STREAMING:
                writer.write(DOUBLE_INDENT + "public StreamObserver<" + method.getJavaInputType() + "> " + 
                           method.getJavaMethodName() + 
                           "(StreamObserver<" + method.getJavaOutputType() + "> responseObserver) {\n");
                writer.write(TRIPLE_INDENT + "return asyncBidiStreamingCall(\n");
                writer.write(TRIPLE_INDENT + INDENT + "getChannel().newCall(" + method.getMethodDescriptorFieldName() + ", getCallOptions()),\n");
                writer.write(TRIPLE_INDENT + INDENT + "responseObserver);\n");
                break;
        }
        
        writer.write(DOUBLE_INDENT + "}\n\n");
    }
    
    /**
     * 生成Future存根
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateFutureStub(ServiceDefinition service, Writer writer) throws IOException {
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Future stub for " + service.getServiceName() + "\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static final class " + service.getFutureStubClassName() + 
                    " extends AbstractStub<" + service.getFutureStubClassName() + "> {\n\n");
        
        // 构造函数
        writer.write(DOUBLE_INDENT + "private " + service.getFutureStubClassName() + 
                    "(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "super(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // build方法
        writer.write(DOUBLE_INDENT + "@Override\n");
        writer.write(DOUBLE_INDENT + "protected " + service.getFutureStubClassName() + 
                    " build(MiniGrpcChannel channel, CallOptions callOptions) {\n");
        writer.write(TRIPLE_INDENT + "return new " + service.getFutureStubClassName() + 
                    "(channel, callOptions);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
        
        // 生成Future方法
        for (MethodDefinition method : service.getMethods()) {
            if (method.isUnary()) {
                generateFutureMethod(method, writer);
            }
        }
        
        writer.write(INDENT + "}\n\n");
    }
    
    /**
     * 生成Future方法
     * 
     * @param method 方法定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateFutureMethod(MethodDefinition method, Writer writer) throws IOException {
        writer.write(DOUBLE_INDENT + "/**\n");
        writer.write(DOUBLE_INDENT + " * " + method.getMethodName() + " future call\n");
        writer.write(DOUBLE_INDENT + " */\n");
        writer.write(DOUBLE_INDENT + "public Future<" + method.getJavaOutputType() + "> " + 
                    method.getJavaMethodName() + "(" + method.getJavaInputType() + " request) {\n");
        writer.write(TRIPLE_INDENT + "return futureUnaryCall(\n");
        writer.write(TRIPLE_INDENT + INDENT + "getChannel().newCall(" + method.getMethodDescriptorFieldName() + ", getCallOptions()), request);\n");
        writer.write(DOUBLE_INDENT + "}\n\n");
    }
    
    /**
     * 生成方法描述符
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateMethodDescriptors(ServiceDefinition service, Writer writer) throws IOException {
        for (MethodDefinition method : service.getMethods()) {
            writer.write(INDENT + "private static final MethodDescriptor<" + 
                        method.getJavaInputType() + ", " + method.getJavaOutputType() + "> " +
                        method.getMethodDescriptorFieldName() + " =\n");
            writer.write(DOUBLE_INDENT + "MethodDescriptor.<" + method.getJavaInputType() + ", " + 
                        method.getJavaOutputType() + ">newBuilder()\n");
            writer.write(TRIPLE_INDENT + ".setType(MethodDescriptor.MethodType." + method.getMethodType() + ")\n");
            writer.write(TRIPLE_INDENT + ".setFullMethodName(generateFullMethodName(SERVICE_NAME, \"" + 
                        method.getMethodName() + "\"))\n");
            writer.write(TRIPLE_INDENT + ".setRequestMarshaller(new JsonMarshaller<>(" + 
                        method.getJavaInputType() + ".class))\n");
            writer.write(TRIPLE_INDENT + ".setResponseMarshaller(new JsonMarshaller<>(" + 
                        method.getJavaOutputType() + ".class))\n");
            writer.write(TRIPLE_INDENT + ".build();\n\n");
        }
    }
    
    /**
     * 生成服务描述符
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateServiceDescriptor(ServiceDefinition service, Writer writer) throws IOException {
        writer.write(INDENT + "private static final ServerServiceDefinition SERVICE_DEFINITION =\n");
        writer.write(DOUBLE_INDENT + "ServerServiceDefinition.builder(SERVICE_NAME)\n");
        
        for (MethodDefinition method : service.getMethods()) {
            writer.write(TRIPLE_INDENT + ".addMethod(" + method.getMethodDescriptorFieldName() + ", /* handler */)\n");
        }
        
        writer.write(TRIPLE_INDENT + ".build();\n\n");
    }
    
    /**
     * 生成工厂方法
     * 
     * @param service 服务定义
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateFactoryMethods(ServiceDefinition service, Writer writer) throws IOException {
        // newBlockingStub
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Creates a new blocking-style stub\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static " + service.getBlockingStubClassName() + 
                    " newBlockingStub(MiniGrpcChannel channel) {\n");
        writer.write(DOUBLE_INDENT + "return new " + service.getBlockingStubClassName() + 
                    "(channel, CallOptions.DEFAULT);\n");
        writer.write(INDENT + "}\n\n");
        
        // newStub
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Creates a new async stub\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static " + service.getAsyncStubClassName() + 
                    " newStub(MiniGrpcChannel channel) {\n");
        writer.write(DOUBLE_INDENT + "return new " + service.getAsyncStubClassName() + 
                    "(channel, CallOptions.DEFAULT);\n");
        writer.write(INDENT + "}\n\n");
        
        // newFutureStub
        writer.write(INDENT + "/**\n");
        writer.write(INDENT + " * Creates a new future stub\n");
        writer.write(INDENT + " */\n");
        writer.write(INDENT + "public static " + service.getFutureStubClassName() + 
                    " newFutureStub(MiniGrpcChannel channel) {\n");
        writer.write(DOUBLE_INDENT + "return new " + service.getFutureStubClassName() + 
                    "(channel, CallOptions.DEFAULT);\n");
        writer.write(INDENT + "}\n\n");
        
        // generateFullMethodName
        writer.write(INDENT + "private static String generateFullMethodName(String serviceName, String methodName) {\n");
        writer.write(DOUBLE_INDENT + "return serviceName + \"/\" + methodName;\n");
        writer.write(INDENT + "}\n\n");
    }
    
    /**
     * 生成文件尾部
     * 
     * @param writer 输出写入器
     * @throws IOException 如果写入失败
     */
    private void generateFileFooter(Writer writer) throws IOException {
        writer.write("}\n");
    }
} 